"""添加分析时段字段并简化任务模型

Revision ID: 03678fe142c6
Revises: 4b8e1a3f9c2d
Create Date: 2025-09-13 10:25:19.439680

"""
from typing import Sequence, Union

from alembic import op
import sqlalchemy as sa
from sqlalchemy.dialects import sqlite

# revision identifiers, used by Alembic.
revision: str = '03678fe142c6'
down_revision: Union[str, Sequence[str], None] = '4b8e1a3f9c2d'
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def upgrade() -> None:
    """Upgrade schema."""
    # ### commands auto generated by Alembic - please adjust! ###
    
    # 尝试删除服务相关的表（如果存在）
    try:
        op.drop_index('ix_service_config_templates_id', table_name='service_config_templates')
    except:
        pass
    
    try:
        op.drop_table('service_config_templates')
    except:
        pass
    
    try:
        op.drop_table('service_classes')
    except:
        pass
    
    try:
        op.drop_table('service_instances')
    except:
        pass
    
    try:
        op.drop_table('service_usage_logs')
    except:
        pass
    
    # SQLite不支持ALTER COLUMN，需要重建表来修改字段类型
    # 备份现有数据
    op.execute("CREATE TABLE analysis_task_backup AS SELECT * FROM analysis_task")
    
    # 删除旧表
    op.drop_table('analysis_task')
    
    # 创建新表，使用新的字段类型
    op.create_table('analysis_task',
        sa.Column('id', sa.Integer(), nullable=False),
        sa.Column('name', sa.String(length=200), nullable=False),
        sa.Column('task_type', sa.String(length=200), nullable=False, comment='任务类型（支持多选，逗号分隔）'),
        sa.Column('execution_type', sa.String(length=50), nullable=False),
        sa.Column('status', sa.String(length=20), nullable=True),
        sa.Column('analysis_start_date', sa.DateTime(), nullable=True, comment='分析开始日期'),
        sa.Column('analysis_end_date', sa.DateTime(), nullable=True, comment='分析结束日期'),
        sa.Column('config', sa.JSON(), nullable=True),
        sa.Column('services', sa.JSON(), nullable=True),
        sa.Column('start_time', sa.DateTime(), nullable=True),
        sa.Column('end_time', sa.DateTime(), nullable=True),
        sa.Column('duration', sa.Float(), nullable=True),
        sa.Column('progress', sa.Float(), nullable=True),
        sa.Column('total_count', sa.Integer(), nullable=True),
        sa.Column('success_count', sa.Integer(), nullable=True),
        sa.Column('failed_count', sa.Integer(), nullable=True),
        sa.Column('retry_count', sa.Integer(), nullable=True),
        sa.Column('max_retries', sa.Integer(), nullable=True),
        sa.Column('data_source_id', sa.Integer(), nullable=True),
        sa.Column('creator_id', sa.Integer(), nullable=True),
        sa.Column('error_message', sa.Text(), nullable=True),
        sa.Column('error_details', sa.JSON(), nullable=True),
        sa.Column('created_at', sa.DateTime(), nullable=True),
        sa.Column('updated_at', sa.DateTime(), nullable=True),
        sa.Column('is_deleted', sa.Boolean(), nullable=True),
        sa.Column('is_active', sa.Boolean(), nullable=True),
        sa.ForeignKeyConstraint(['creator_id'], ['user.id'], ),
        sa.ForeignKeyConstraint(['data_source_id'], ['data_source.id'], ),
        sa.PrimaryKeyConstraint('id')
    )
    
    # 从备份表恢复数据（排除description字段）
    op.execute("""
        INSERT INTO analysis_task (
            id, name, task_type, execution_type, status, config, services, 
            start_time, end_time, duration, progress, total_count, success_count, 
            failed_count, retry_count, max_retries, data_source_id, creator_id, 
            error_message, error_details, created_at, updated_at, is_deleted, is_active,
            analysis_start_date, analysis_end_date
        )
        SELECT 
            id, name, task_type, execution_type, status, config, services, 
            start_time, end_time, duration, progress, total_count, success_count, 
            failed_count, retry_count, max_retries, data_source_id, creator_id, 
            error_message, error_details, created_at, updated_at, is_deleted, is_active,
            NULL, NULL
        FROM analysis_task_backup
    """)
    
    # 删除备份表
    op.drop_table('analysis_task_backup')
    # ### end Alembic commands ###


def downgrade() -> None:
    """Downgrade schema."""
    # ### commands auto generated by Alembic - please adjust! ###
    
    # SQLite不支持ALTER COLUMN，需要重建表来恢复原字段类型
    # 备份现有数据
    op.execute("CREATE TABLE analysis_task_backup AS SELECT * FROM analysis_task")
    
    # 删除当前表
    op.drop_table('analysis_task')
    
    # 重建原始表结构（包含description字段）
    op.create_table('analysis_task',
        sa.Column('id', sa.Integer(), nullable=False),
        sa.Column('name', sa.String(length=200), nullable=False),
        sa.Column('task_type', sa.String(length=50), nullable=False),
        sa.Column('execution_type', sa.String(length=50), nullable=False),
        sa.Column('status', sa.String(length=20), nullable=True),
        sa.Column('config', sa.JSON(), nullable=True),
        sa.Column('services', sa.JSON(), nullable=True),
        sa.Column('start_time', sa.DateTime(), nullable=True),
        sa.Column('end_time', sa.DateTime(), nullable=True),
        sa.Column('duration', sa.Float(), nullable=True),
        sa.Column('progress', sa.Float(), nullable=True),
        sa.Column('total_count', sa.Integer(), nullable=True),
        sa.Column('success_count', sa.Integer(), nullable=True),
        sa.Column('failed_count', sa.Integer(), nullable=True),
        sa.Column('retry_count', sa.Integer(), nullable=True),
        sa.Column('max_retries', sa.Integer(), nullable=True),
        sa.Column('data_source_id', sa.Integer(), nullable=True),
        sa.Column('creator_id', sa.Integer(), nullable=True),
        sa.Column('error_message', sa.Text(), nullable=True),
        sa.Column('error_details', sa.JSON(), nullable=True),
        sa.Column('created_at', sa.DateTime(), nullable=True),
        sa.Column('updated_at', sa.DateTime(), nullable=True),
        sa.Column('is_deleted', sa.Boolean(), nullable=True),
        sa.Column('is_active', sa.Boolean(), nullable=True),
        sa.Column('description', sa.Text(), nullable=True),
        sa.ForeignKeyConstraint(['creator_id'], ['user.id'], ),
        sa.ForeignKeyConstraint(['data_source_id'], ['data_source.id'], ),
        sa.PrimaryKeyConstraint('id')
    )
    
    # 从备份表恢复数据（添加description字段）
    op.execute("""
        INSERT INTO analysis_task (
            id, name, task_type, execution_type, status, config, services, 
            start_time, end_time, duration, progress, total_count, success_count, 
            failed_count, retry_count, max_retries, data_source_id, creator_id, 
            error_message, error_details, created_at, updated_at, is_deleted, is_active,
            description
        )
        SELECT 
            id, name, task_type, execution_type, status, config, services, 
            start_time, end_time, duration, progress, total_count, success_count, 
            failed_count, retry_count, max_retries, data_source_id, creator_id, 
            error_message, error_details, created_at, updated_at, is_deleted, is_active,
            NULL
        FROM analysis_task_backup
    """)
    
    # 删除备份表
    op.drop_table('analysis_task_backup')
    
    # 重建被删除的服务相关表
    op.create_table('service_usage_logs',
    sa.Column('id', sa.INTEGER(), nullable=False),
    sa.Column('service_class_id', sa.INTEGER(), nullable=False),
    sa.Column('service_instance_id', sa.INTEGER(), nullable=True),
    sa.Column('user_id', sa.INTEGER(), nullable=True),
    sa.Column('request_data', sqlite.JSON(), nullable=True),
    sa.Column('response_data', sqlite.JSON(), nullable=True),
    sa.Column('execution_time', sa.INTEGER(), nullable=True),
    sa.Column('token_count', sa.INTEGER(), nullable=True),
    sa.Column('cost', sa.VARCHAR(length=50), nullable=True),
    sa.Column('status', sa.VARCHAR(length=20), nullable=True),
    sa.Column('error_message', sa.TEXT(), nullable=True),
    sa.Column('created_at', sa.DATETIME(), nullable=True),
    sa.Column('updated_at', sa.DATETIME(), nullable=True),
    sa.Column('is_deleted', sa.BOOLEAN(), nullable=True),
    sa.ForeignKeyConstraint(['service_class_id'], ['service_classes.id'], ),
    sa.ForeignKeyConstraint(['service_instance_id'], ['service_instances.id'], ),
    sa.ForeignKeyConstraint(['user_id'], ['user.id'], ),
    sa.PrimaryKeyConstraint('id')
    )
    op.create_table('service_instances',
    sa.Column('id', sa.INTEGER(), nullable=False),
    sa.Column('service_class_id', sa.INTEGER(), nullable=False),
    sa.Column('user_id', sa.INTEGER(), nullable=True),
    sa.Column('instance_name', sa.VARCHAR(length=100), nullable=False),
    sa.Column('instance_config', sqlite.JSON(), nullable=True),
    sa.Column('is_active', sa.BOOLEAN(), nullable=True),
    sa.Column('status', sa.VARCHAR(length=20), nullable=True),
    sa.Column('usage_count', sa.INTEGER(), nullable=True),
    sa.Column('last_used_at', sa.DATETIME(), nullable=True),
    sa.Column('created_at', sa.DATETIME(), nullable=True),
    sa.Column('updated_at', sa.DATETIME(), nullable=True),
    sa.Column('is_deleted', sa.BOOLEAN(), nullable=True),
    sa.ForeignKeyConstraint(['service_class_id'], ['service_classes.id'], ),
    sa.ForeignKeyConstraint(['user_id'], ['user.id'], ),
    sa.PrimaryKeyConstraint('id')
    )
    op.create_table('service_classes',
    sa.Column('id', sa.INTEGER(), nullable=False),
    sa.Column('service_id', sa.VARCHAR(length=50), nullable=False),
    sa.Column('name', sa.VARCHAR(length=100), nullable=False),
    sa.Column('description', sa.TEXT(), nullable=True),
    sa.Column('category', sa.VARCHAR(length=20), nullable=True),
    sa.Column('provider', sa.VARCHAR(length=20), nullable=True),
    sa.Column('model_code', sa.VARCHAR(length=100), nullable=True),
    sa.Column('max_tokens', sa.INTEGER(), nullable=True),
    sa.Column('temperature', sa.FLOAT(), nullable=True),
    sa.Column('prompt_template_id', sa.INTEGER(), nullable=True),
    sa.Column('is_active', sa.BOOLEAN(), nullable=True),
    sa.Column('config', sqlite.JSON(), nullable=True),
    sa.Column('usage_count', sa.INTEGER(), nullable=True),
    sa.Column('last_used_at', sa.DATETIME(), nullable=True),
    sa.Column('status', sa.VARCHAR(length=20), nullable=True),
    sa.Column('error_message', sa.TEXT(), nullable=True),
    sa.Column('created_at', sa.DATETIME(), nullable=True),
    sa.Column('updated_at', sa.DATETIME(), nullable=True),
    sa.Column('is_deleted', sa.BOOLEAN(), nullable=True),
    sa.ForeignKeyConstraint(['prompt_template_id'], ['prompt_templates.id'], ),
    sa.PrimaryKeyConstraint('id'),
    sa.UniqueConstraint('service_id')
    )
    op.create_table('service_config_templates',
    sa.Column('name', sa.VARCHAR(length=100), nullable=False),
    sa.Column('description', sa.TEXT(), nullable=True),
    sa.Column('service_category', sa.VARCHAR(length=20), nullable=True),
    sa.Column('config_schema', sqlite.JSON(), nullable=True),
    sa.Column('default_config', sqlite.JSON(), nullable=True),
    sa.Column('usage_count', sa.INTEGER(), nullable=True),
    sa.Column('created_at', sa.DATETIME(), nullable=True),
    sa.Column('updated_at', sa.DATETIME(), nullable=True),
    sa.Column('id', sa.INTEGER(), nullable=False),
    sa.Column('is_deleted', sa.BOOLEAN(), nullable=True),
    sa.PrimaryKeyConstraint('id')
    )
    op.create_index(op.f('ix_service_config_templates_id'), 'service_config_templates', ['id'], unique=False)
    # ### end Alembic commands ###
